from futu import *
# from futu_daily_pick import picked_codes
import numpy as np
import logging
from time import sleep
from snb_trade import snowball
from utility import code_futu_to_snowball
# 设置日志记录器
logging.basicConfig(filename='log.log', level=logging.INFO, format='%(asctime)s - %(message)s')
snb = snowball()
picking_codes = ['HK.02202','HK.06030']
picked_codes = code_futu_to_snowball(picking_codes)
quote_ctx = OpenQuoteContext(host='127.0.0.1', port=11111)

def subscribe(futu_code):
    ret_sub, err_message = quote_ctx.subscribe([futu_code], [SubType.K_DAY,SubType.K_60M,SubType.TICKER,SubType.ORDER_BOOK], subscribe_push=False)
    # 先订阅 K 线类型。订阅成功后 OpenD 将持续收到服务器的推送，False 代表暂时不需要推送给脚本
    if ret_sub == RET_OK:
        return True
    else:
        print('订阅失败', err_message)
    # 关闭当条连接，OpenD 会在1分钟后自动取消相应股票相应类型的订阅
    quote_ctx.close()  

def is_normal_trading_time():
    ret, data = quote_ctx.get_market_state(picking_codes)  
    if ret != RET_OK:
        print('获取市场状态失败：', data)
        return False
    market_state = data['market_state'][0]
    '''
    MarketState.MORNING            港、A 股早盘
    MarketState.AFTERNOON          港、A 股下午盘，美股全天
    MarketState.FUTURE_DAY_OPEN    港、新、日期货日市开盘
    MarketState.FUTURE_OPEN        美期货开盘
    MarketState.FUTURE_BREAK_OVER  美期货休息后开盘
    MarketState.NIGHT_OPEN         港、新、日期货夜市开盘
    '''
    if  market_state == MarketState.MORNING or \
        market_state == MarketState.AFTERNOON or \
        market_state == MarketState.FUTURE_DAY_OPEN  or \
        market_state == MarketState.FUTURE_OPEN  or \
        market_state == MarketState.FUTURE_BREAK_OVER  or \
        market_state == MarketState.NIGHT_OPEN:
        return True
    return False

def get_kline(futu_code):
    # if subscribe(futu_code):
        ret, data = quote_ctx.get_cur_kline(futu_code, 2, KLType.K_60M, AuType.QFQ) 
        if ret == RET_OK:
            print(data)
            print(data['turnover_rate'][0])   # 取第一条的换手率
            print(data['turnover_rate'].values.tolist())   # 转为 list
        else:
            print('获取K线error:', data)

def ticker(futu_code):
    # if subscribe(futu_code):
        # 获取最近2个逐笔
        ret, data = quote_ctx.get_rt_ticker(futu_code, 2)  
        if ret != RET_OK or data.empty:
            print('逐笔error:', data)
            return None
        return data['price'][0],data['volume'][0],data['turnover'][0],data['ticker_direction'][0]

def data_order_book(futu_code):
    ret, data = quote_ctx.get_order_book(futu_code, num=1)
    if ret != RET_OK:
        print('order_book_error:', data)
        return None, None, None, None
    return data['Ask'][0][0],data['Ask'][0][1],data['Bid'][0][0],data['Bid'][0][1]


############################ 填充以下函数来完成您的策略 ############################
def on_init():
    print('************ 开始运行 ***********')
    return True

# 每个 tick 运行一次，可将策略的主要逻辑写在此处
def on_tick():
    if not is_normal_trading_time():
        print('not trading time')
        # return
    holdings = snb.holdings()
    num_holding = 10
    leg_value = snb.net_value()/num_holding
    if holdings ==0 or holdings.empty:
        print(f'空仓,杠杆率: {snb.leverage()}')
    else:
        for code in holdings['symbol'].tolist():
            pos = holdings[holdings['symbol'] == code]['position'].iloc[0]
            if code not in picked_codes:
                snb.sell(code,pos)
                print(f'清仓:{code}')
    dict = {}
    for code in picking_codes:
        sleep(3)
        if code not in picked_codes:
            dict[code] = pd.DataFrame({
                        'code':code,
                        'last_price': ticker(code)[0],
                        'vol': ticker(code)[1],
                        'turnover': ticker(code)[2],
                        'direction': ticker(code)[3],
                        'ask_price': data_order_book(code)[0],
                        # 'ask_vol': data_order_book(code)[1],
                        'bid_price': data_order_book(code)[2],
                        # 'bid_vol': data_order_book(code)[3],
                        },index=[0])
        else:
            dict[code].loc[len(dict[code].index)] = {
                        'code':code,
                        'last_price': ticker(code)[0],
                        'vol': ticker(code)[1],
                        'turnover': ticker(code)[2],
                        'direction': ticker(code)[3],
                        'ask_price': data_order_book(code)[0],
                        # 'ask_vol': data_order_book(code)[1],
                        'bid_price': data_order_book(code)[2],
                        # 'bid_vol': data_order_book(code)[3],
                        }
        # 条件
        dict[code]['unit_bid'] = 0
        dict[code]['unit_ask'] = 0
        dict[code]['cum_ratio'] = 0
        if dict[code]['direction'].iloc[-1] == 'BUY':
            # data[code]['unit_bid'] = data[code]['vol'] + data[code]['bid_vol']
            dict[code]['unit_bid'] = dict[code]['turnover'].iloc[-1]
        elif dict[code]['direction'].iloc[-1] =='SELL':
            # data[code]['unit_ask'] = data[code]['vol'] + data[code]['ask_vol']
            dict[code]['unit_ask'] = dict[code]['turnover'].iloc[-1]
        elif dict[code]['direction'].iloc[-1] == 'NEUTRAL' and dict[code].shape[0] > 1:
            if dict[code]['ask_price'].iloc[-1] == dict[code]['ask_price'].iloc[-2]:
                # data[code]['unit_bid'] = data[code]['bid_vol'].shift(1) - data[code]['bid_vol']
                dict[code]['unit_bid'] = dict[code]['turnover'].iloc[-1]
            elif dict[code]['bid_price'].iloc[-1] == dict[code]['bid_price'].iloc[-2]:
                # data[code]['unit_ask'] = data[code]['ask_vol'].shift(1) - data[code]['ask_vol']
                dict[code]['unit_ask'] = dict[code]['turnover'].iloc[-1]
        dict[code]['unit'] = dict[code]['unit_bid'] - dict[code]['unit_ask']
        if dict[code]['unit'].iloc[-1] > 10000:
        # sgn = unit / max(abs(unit), 1)
        # cum_ratio += np.log(unit +1)*sgn
            dict[code]['cum_ratio'] += dict[code]['unit']
            logging.info(f'code:{code},cum_ratio:{dict[code]["cum_ratio"].iloc[-1]}')
            # print(dict[code])
        if not dict[code].empty: 
            # 数量和价格
            leg_vol = 100
            if dict[code]['last_price'].iloc[-1] < 100:
                leg_vol = int(max(int(leg_value/dict[code]['last_price'].iloc[-1]/1000), 1)*1000)
            # elif dict[code]['last_price'].iloc[-1] < 200:
            #     leg_vol = 500
            code_snb = code[3:]
            if dict[code]['cum_ratio'].iloc[-1] > 20000:
                if holdings == 0 or holdings.empty:
                    # snb.buy(code_snb,leg_vol)
                    print(f'建仓:{code_snb},数量:{leg_vol}')
                elif leg_vol > holdings[holdings['symbol']==code_snb]['position'].iloc[0]:
                    # snb.buy(code_snb,leg_vol-holdings[holdings['symbol']==code_snb]['position'][0])
                    print(f'加仓:{code_snb}')
            elif dict[code]['cum_ratio'].iloc[-1] < -20000:
                if leg_vol < holdings[holdings['symbol']==code_snb]['position'].iloc[0]:
                    # snb.sell(code_snb,holdings[holdings['symbol']==code_snb]['position'][0]-leg_vol)
                    print(f'减仓:{code_snb}')
        
def on_bar():
    for code in picking_codes:
        print(get_kline(code))


################################ 框架实现部分，可忽略不看 ###############################
class OnTickClass(TickerHandlerBase):
    def on_recv_rsp(self, rsp_pb):
        on_tick()


class OnBarClass(CurKlineHandlerBase):
    last_time = None
    def on_recv_rsp(self, rsp_pb):
        ret_code, data = super(OnBarClass, self).on_recv_rsp(rsp_pb)
        if ret_code == RET_OK:
            cur_time = data['time_key'][0]
            if cur_time != self.last_time and data['k_type'][0] == KLType.K_1M:
                if self.last_time is not None:
                    on_bar()
                self.last_time = cur_time


# 主函数
if __name__ == '__main__':
    # 初始化策略
    if not on_init():
        print('策略初始化失败，脚本退出！')
        quote_ctx.close()
        # trade_context.close()
    else:
        # 设置回调
        quote_ctx.set_handler(OnTickClass())
        quote_ctx.set_handler(OnBarClass())
        # 订阅标的合约的 逐笔，K 线和摆盘，以便获取数据
        quote_ctx.subscribe(code_list=picking_codes, subtype_list=[SubType.TICKER, SubType.ORDER_BOOK, KLType.K_60M,KLType.K_DAY])

